from six import StringIO
from sklearn import model_selection, metrics
from sklearn.tree import DecisionTreeClassifier
from sklearn.tree import export_graphviz
import pandas as pd
import numpy as np
import pydotplus


class PredictDiabetes(object):
    @staticmethod
    def predict_diabetes():
        ds = pd.read_csv('pima-indians-diabetes.csv')
        feature_columns = ['pregnant', 'glucose', 'bp', 'skin', 'insulin', 'bmi', 'pedigree', 'age']
        X = ds[feature_columns]
        y = ds.label
        print(X, y)

        dtc = DecisionTreeClassifier()
        dtc.fit(X, y)

        test_data = [8, 125, 96, 0, 0, 0, 0.232, 24]
        y_pred = dtc.predict(np.array(test_data).reshape(1, -1))
        print(y_pred)

    @staticmethod
    def predict_diabetes_graph():
        ds = pd.read_csv('pima-indians-diabetes.csv')
        feature_columns = ['pregnant', 'glucose', 'bp', 'skin', 'insulin', 'bmi', 'pedigree', 'age']
        X = ds[feature_columns]
        y = ds.label
        # print(X, y)

        dtc = DecisionTreeClassifier(criterion='gini')
        dtc = DecisionTreeClassifier(criterion='entropy')
        dtc.fit(X, y)

        test_data = [8, 125, 96, 0, 0, 0, 0.232, 24]
        y_pred = dtc.predict(np.array(test_data).reshape(1, -1))
        print(y_pred)

        dot_data = StringIO()
        export_graphviz(dtc, out_file=dot_data,
                        feature_names=feature_columns,
                        class_names=['0', '1'],
                        rounded=True,
                        filled=True,
                        special_characters=True)
        graph = pydotplus.graph_from_dot_data(dot_data.getvalue())
        graph.write_pdf('decision_tree.pdf')

    @staticmethod
    def get_accuracy_score(test_size=0.25):
        ds = pd.read_csv('pima-indians-diabetes.csv')
        feature_columns = ['pregnant', 'glucose', 'bp', 'skin', 'insulin', 'bmi', 'pedigree', 'age']
        X = ds[feature_columns]
        y = ds.label
        X_train, X_test, y_train, y_test = model_selection.train_test_split(X, y, test_size=test_size)
        dtc = DecisionTreeClassifier()
        dtc.fit(X_train, y_train)
        y_pred = dtc.predict(X_test)
        score = metrics.accuracy_score(y_test, y_pred)
        print(score)


if __name__ == '__main__':
    # PredictDiabetes.predict_diabetes()
    # PredictDiabetes.get_accuracy_score(0.1)
    PredictDiabetes.predict_diabetes_graph()
